import cv2
import numpy as np

# 定义颜色及其HSV范围（黑色需特殊处理）
color_config = {
    "black": {"lower": [0, 0, 0], "upper": [180, 255, 50]},  # 低明度
    "yellow": {"lower": [20, 100, 100], "upper": [30, 255, 255]},  # 黄色
    "red": {"lower": [0, 100, 100], "upper": [10, 255, 255]},  # 红色（需双区间）
    "blue": {"lower": [90, 50, 50], "upper": [130, 255, 255]}  # 蓝色
}


def detect_colored_circles(image):
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    result = image.copy()

    for color_name, config in color_config.items():
        # 红色需要处理两个色相区间
        if color_name == "red":
            lower1 = np.array(config["lower"], dtype=np.uint8)
            upper1 = np.array(config["upper"], dtype=np.uint8)
            lower2 = np.array([170, config["lower"][1], config["lower"][2]], dtype=np.uint8)
            upper2 = np.array([180, config["upper"][1], config["upper"][2]], dtype=np.uint8)
            mask1 = cv2.inRange(hsv, lower1, upper1)
            mask2 = cv2.inRange(hsv, lower2, upper2)
            mask = cv2.bitwise_or(mask1, mask2)
        else:
            lower = np.array(config["lower"], dtype=np.uint8)
            upper = np.array(config["upper"], dtype=np.uint8)
            mask = cv2.inRange(hsv, lower, upper)

        # 形态学操作（去噪+填充）
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)  # 开运算去噪
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)  # 闭运算填充

        cv2.imshow(color_name,mask)
        """        # 方法1：霍夫圆变换检测
        circles = cv2.HoughCircles(
            mask, cv2.HOUGH_GRADIENT, dp=1.2, minDist=30,
            param1=50, param2=30, minRadius=10, maxRadius=100
        )

        if circles is not None:
            circles = np.uint16(np.around(circles))
            for circle in circles[0, :]:
                x, y, r = circle[0], circle[1], circle[2]
                cv2.circle(result, (x, y), r, (0, 255, 0), 3)  # 画圆
                cv2.putText(result, color_name, (x - r, y - r),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)"""

        # 方法2（备选）：轮廓分析（圆度筛选）
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area < 100: continue  # 过滤小区域
            perimeter = cv2.arcLength(cnt, True)
            if perimeter == 0: continue
            circularity = 4 * np.pi * area / (perimeter ** 2)  # 圆度计算
            if circularity > 0.7:  # 圆度阈值（1为完美圆）
                (x, y), r = cv2.minEnclosingCircle(cnt)
                cv2.circle(result, (int(x), int(y)), int(r), (0, 255, 0), 2)
                cv2.putText(result, color_name, (int(x) - 20, int(y) - 20),
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

    return result

capture = cv2.VideoCapture(1)  # 打开笔记本内置摄像头
while (capture.isOpened()):  # 笔记本内置摄像头被打开后
    retval, image = capture.read(1)  # 从摄像头中实时读取视频
    # img用做绘图，image用作处理
    img = image.copy()


    result_image = detect_colored_circles(image)

    cv2.imshow("Video", result_image)  # 在窗口中显示读取到的视频

    key = cv2.waitKey(1)  # 窗口的图像刷新时间为1毫秒
    if key == 27:  # 如果按下esc
        break
capture.release()  # 关闭笔记本内置摄像头
cv2.destroyAllWindows()  # 销毁显示摄像头视频的窗口